from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import matplotlib.pyplot as plt
import time
import svgwrite
from datetime import datetime
import matplotlib.cm as cm
import matplotlib.colors as mcolors

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json"):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        print("state: ", state)
        print("goal: ", goal)
        print("landscape: ", landscape)
    return landscape, state, goal

# save current landscape as json file
def save_current_task(landscape, start_x, start_y, goal_x, goal_y, pth = "./logs/landscape.json"):
    landscape_ = []
    for j in range(landscape[0].shape[0]):
        landscape_.append(int(landscape[0][j]))

    with open(pth, "w") as f:
        json.dump({"data": landscape_, 
                   "state": [start_x, start_y],
                   "goal": [goal_x, goal_y]}, f)

def visualize_hidden_state(hidden_state, window_name="Neural Hidden State"):
    """
    可视化神经网络隐状态 - 用简洁配色的实心圆横向排列
    
    Args:
        hidden_state: numpy数组，包含隐状态的数值
        window_name: 窗口名称
    """
    if hidden_state is None or len(hidden_state) == 0:
        return
    
    # 转换JAX数组为numpy数组（如果需要的话）
    if hasattr(hidden_state, 'tolist'):
        hidden_values = np.array(hidden_state.tolist())
    else:
        hidden_values = np.array(hidden_state)
    
    # 确保是1D数组
    if hidden_values.ndim > 1:
        hidden_values = hidden_values.flatten()
    
    original_num_units = len(hidden_values)
    if original_num_units == 0:
        return
    
    # 等间隔选取32个单元进行可视化
    target_units = 32
    if original_num_units > target_units:
        indices = np.linspace(0, original_num_units - 1, target_units, dtype=int)
        hidden_values = hidden_values[indices]
        num_units = target_units
    else:
        num_units = original_num_units
    
    # 计算可视化参数 - 确保圆圈不重叠
    circle_radius = 8  # 减小圆圈半径
    circle_diameter = circle_radius * 2
    circle_spacing = circle_diameter + 6  # 圆圈间距确保不重叠
    margin = 20  # 减小边距
    
    # 计算画布尺寸
    canvas_width = 2 * margin + (num_units - 1) * circle_spacing + circle_diameter
    canvas_height = 2 * margin + circle_diameter
    
    # 创建画布 - 简洁的深灰色背景
    background_color = (35, 35, 40)  # 深灰色，简洁而不廉价
    img = np.full((canvas_height, canvas_width, 3), background_color, dtype=np.uint8)
    
    # 规范化隐状态值到[0,1]范围
    min_val = np.min(hidden_values)
    max_val = np.max(hidden_values)
    
    if max_val > min_val:
        normalized_values = (hidden_values - min_val) / (max_val - min_val)
    else:
        normalized_values = np.ones_like(hidden_values) * 0.5
    
    # 绘制每个隐状态单元 - 简洁配色方案 + 可变半径
    for i, norm_val in enumerate(normalized_values):
        # 计算圆心位置
        center_x = margin + i * circle_spacing + circle_radius
        center_y = canvas_height // 2
        
        # 简洁的蓝色渐变配色方案：从深蓝到亮蓝
        # 基础颜色：深蓝(40,60,100) -> 亮蓝(120,160,220)
        blue_intensity = int(40 + norm_val * 80)   # 40 -> 120
        green_intensity = int(60 + norm_val * 100) # 60 -> 160  
        red_intensity = int(100 + norm_val * 120)  # 100 -> 220
        
        main_color = (blue_intensity, green_intensity, red_intensity)
        
        # 根据数值大小计算圆圈半径 - 数值越小半径越小
        min_radius = 3  # 最小半径
        max_radius = circle_radius  # 最大半径
        current_radius = int(min_radius + norm_val * (max_radius - min_radius))
        
        # 绘制主体圆圈（反锯齿，可变半径）
        cv2.circle(img, (center_x, center_y), current_radius, main_color, -1, cv2.LINE_AA)
    
    # 显示图像
    cv2.imshow(window_name, img)

def create_single_neural_state_svg_rotated(hidden_state, width=60, height=800):
    """
    创建单个神经状态的SVG表示，旋转90度（垂直排列圆点）
    
    Args:
        hidden_state: numpy数组，包含隐状态的数值
        width: SVG宽度
        height: SVG高度
    
    Returns:
        svgwrite.Drawing 对象
    """
    if hidden_state is None or len(hidden_state) == 0:
        return None
    
    # 转换JAX数组为numpy数组（如果需要的话）
    if hasattr(hidden_state, 'tolist'):
        hidden_values = np.array(hidden_state.tolist())
    else:
        hidden_values = np.array(hidden_state)
    
    # 确保是1D数组
    if hidden_values.ndim > 1:
        hidden_values = hidden_values.flatten()
    
    original_num_units = len(hidden_values)
    if original_num_units == 0:
        return None
    
    # 等间隔选取32个单元进行可视化
    target_units = 32
    if original_num_units > target_units:
        indices = np.linspace(0, original_num_units - 1, target_units, dtype=int)
        hidden_values = hidden_values[indices]
        num_units = target_units
    else:
        num_units = original_num_units
    
    # 计算可视化参数 - 旋转90度，所以现在是垂直排列
    circle_radius = 8  # 圆圈半径
    circle_diameter = circle_radius * 2
    circle_spacing = circle_diameter + 6  # 圆圈间距
    margin = 20  # 边距
    
    # 创建SVG画布
    dwg = svgwrite.Drawing(size=(width, height))
    
    # 背景
    dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='rgb(35,35,40)'))
    
    # 规范化隐状态值到[0,1]范围
    min_val = np.min(hidden_values)
    max_val = np.max(hidden_values)
    
    if max_val > min_val:
        normalized_values = (hidden_values - min_val) / (max_val - min_val)
    else:
        normalized_values = np.ones_like(hidden_values) * 0.5
    
    # 绘制每个隐状态单元 - 垂直排列
    for i, norm_val in enumerate(normalized_values):
        # 计算圆心位置（垂直排列）
        center_x = width // 2
        center_y = margin + i * circle_spacing + circle_radius
        
        # 简洁的蓝色渐变配色方案：从深蓝到亮蓝
        blue_intensity = int(40 + norm_val * 80)   # 40 -> 120
        green_intensity = int(60 + norm_val * 100) # 60 -> 160  
        red_intensity = int(100 + norm_val * 120)  # 100 -> 220
        
        # 根据数值大小计算圆圈半径
        min_radius = 3
        max_radius = circle_radius
        current_radius = int(min_radius + norm_val * (max_radius - min_radius))
        
        # 绘制主体圆圈
        dwg.add(dwg.circle(center=(center_x, center_y), r=current_radius, 
                          fill=f'rgb({blue_intensity},{green_intensity},{red_intensity})'))
    
    return dwg

def create_neural_states_sequence_opencv(neural_states_sequence):
    """
    创建神经状态序列的OpenCV图像用于预览
    
    Args:
        neural_states_sequence: 神经状态序列列表
    
    Returns:
        OpenCV图像
    """
    if not neural_states_sequence or len(neural_states_sequence) == 0:
        return None
    
    # 计算实际需要的尺寸
    target_units = 32
    circle_radius = 8
    circle_spacing = circle_radius * 2 + 6  # 22像素
    margin = 10  # 减小顶部和底部边距
    
    # 计算实际高度：顶部边距 + 圆点区域 + 底部边距
    actual_height = margin + (target_units - 1) * circle_spacing + circle_radius * 2 + margin
    
    # 单个神经状态图的尺寸
    single_width = 60
    single_height = actual_height
    separator_width = 4  # 分隔符宽度（像素间隔）
    
    # 计算总的图像尺寸
    total_width = len(neural_states_sequence) * single_width + (len(neural_states_sequence) - 1) * separator_width
    total_height = single_height
    
    # 创建主画布
    background_color = (35, 35, 40)  # 深灰色背景
    img = np.full((total_height, total_width, 3), background_color, dtype=np.uint8)
    
    # 为每个神经状态创建图像并拼接
    current_x = 0
    for i, neural_state in enumerate(neural_states_sequence):
        # 转换神经状态数据
        if hasattr(neural_state, 'tolist'):
            hidden_values = np.array(neural_state.tolist())
        else:
            hidden_values = np.array(neural_state)
        
        if hidden_values.ndim > 1:
            hidden_values = hidden_values.flatten()
        
        original_num_units = len(hidden_values)
        if original_num_units > 0:
            # 选取32个单元
            if original_num_units > target_units:
                indices = np.linspace(0, original_num_units - 1, target_units, dtype=int)
                hidden_values = hidden_values[indices]
                num_units = target_units
            else:
                num_units = original_num_units
            
            # 规范化
            min_val = np.min(hidden_values)
            max_val = np.max(hidden_values)
            
            if max_val > min_val:
                normalized_values = (hidden_values - min_val) / (max_val - min_val)
            else:
                normalized_values = np.ones_like(hidden_values) * 0.5
            
            # 绘制单个神经状态的背景矩形（垂直排列）
            cv2.rectangle(img, (current_x, 0), (current_x + single_width, single_height), (45, 45, 50), -1)
            
            # 绘制圆点（垂直排列）
            for j, norm_val in enumerate(normalized_values):
                center_x = current_x + single_width // 2
                center_y = margin + j * circle_spacing + circle_radius
                
                # 颜色计算 - 与OpenCV原版保持一致
                blue_intensity = int(40 + norm_val * 80)   # 40 -> 120
                green_intensity = int(60 + norm_val * 100) # 60 -> 160  
                red_intensity = int(100 + norm_val * 120)  # 100 -> 220
                
                main_color = (blue_intensity, green_intensity, red_intensity)  # BGR格式
                
                # 半径计算
                min_radius = 3
                max_radius = circle_radius
                current_radius = int(min_radius + norm_val * (max_radius - min_radius))
                
                # 绘制圆圈
                cv2.circle(img, (center_x, center_y), current_radius, main_color, -1, cv2.LINE_AA)
        
        # 更新x位置
        current_x += single_width + separator_width
    
    return img

def save_neural_states_sequence_svg(neural_states_sequence, pth="./logs/neural_states_sequence.svg"):
    """
    将神经状态序列横向拼接保存为SVG
    
    Args:
        neural_states_sequence: 神经状态序列列表
        pth: 保存路径
    """
    if not neural_states_sequence or len(neural_states_sequence) == 0:
        return
    
    # 计算实际需要的尺寸 - 与OpenCV版本保持一致
    target_units = 32
    circle_radius = 8
    circle_spacing = circle_radius * 2 + 6  # 22像素
    margin = 10  # 减小顶部和底部边距
    
    # 计算实际高度：顶部边距 + 圆点区域 + 底部边距
    actual_height = margin + (target_units - 1) * circle_spacing + circle_radius * 2 + margin
    
    # 单个神经状态图的尺寸
    single_width = 60
    single_height = actual_height
    separator_width = 4  # 分隔符宽度（像素间隔）
    
    # 计算总的SVG尺寸
    total_width = len(neural_states_sequence) * single_width + (len(neural_states_sequence) - 1) * separator_width
    total_height = single_height
    
    # 创建主SVG画布
    main_dwg = svgwrite.Drawing(pth, size=(total_width, total_height))
    
    # 背景
    main_dwg.add(main_dwg.rect(insert=(0, 0), size=(total_width, total_height), fill='rgb(35,35,40)'))
    
    # 为每个神经状态创建SVG并拼接
    current_x = 0
    for i, neural_state in enumerate(neural_states_sequence):
        # 添加单个神经状态的背景矩形
        main_dwg.add(main_dwg.rect(insert=(current_x, 0), size=(single_width, single_height), fill='rgb(45,45,50)'))
        
        # 转换神经状态数据
        if hasattr(neural_state, 'tolist'):
            hidden_values = np.array(neural_state.tolist())
        else:
            hidden_values = np.array(neural_state)
        
        if hidden_values.ndim > 1:
            hidden_values = hidden_values.flatten()
        
        original_num_units = len(hidden_values)
        if original_num_units > 0:
            # 选取32个单元
            if original_num_units > target_units:
                indices = np.linspace(0, original_num_units - 1, target_units, dtype=int)
                hidden_values = hidden_values[indices]
                num_units = target_units
            else:
                num_units = original_num_units
            
            # 规范化
            min_val = np.min(hidden_values)
            max_val = np.max(hidden_values)
            
            if max_val > min_val:
                normalized_values = (hidden_values - min_val) / (max_val - min_val)
            else:
                normalized_values = np.ones_like(hidden_values) * 0.5
            
            # 绘制圆点（垂直排列）
            for j, norm_val in enumerate(normalized_values):
                center_x = current_x + single_width // 2
                center_y = margin + j * circle_spacing + circle_radius
                
                # 颜色计算 - 注意BGR转RGB
                blue_intensity = int(40 + norm_val * 80)   # 40 -> 120
                green_intensity = int(60 + norm_val * 100) # 60 -> 160  
                red_intensity = int(100 + norm_val * 120)  # 100 -> 220
                
                # SVG使用RGB格式，所以需要转换BGR -> RGB
                svg_color = f'rgb({red_intensity},{green_intensity},{blue_intensity})'
                
                # 半径计算
                min_radius = 3
                max_radius = circle_radius
                current_radius = int(min_radius + norm_val * (max_radius - min_radius))
                
                # 添加圆圈
                main_dwg.add(main_dwg.circle(center=(center_x, center_y), r=current_radius,
                                           fill=svg_color))
        
        # 更新x位置
        current_x += single_width + separator_width
    
    # 保存SVG
    main_dwg.save()
    print(f"Neural states sequence SVG saved to {pth}")

def save_svg_snapshot(grid, trajectory, init_pos, goal_pos, agent_pos, agent_action, perturbation_arrows, pth):
    """将当前画面保存为高质量SVG矢量图"""
    grid_size_display = 40
    width, height = grid.shape[0], grid.shape[1]
    svg_width = width * grid_size_display
    svg_height = height * grid_size_display
    
    # 创建SVG画布
    dwg = svgwrite.Drawing(pth, size=(svg_width, svg_height))
    
    # 计算agent位置用于可视范围判断
    agent_grid_x = int((agent_pos[0] - 20) // 40)
    agent_grid_y = int((agent_pos[1] - 20) // 40)
    
    # 绘制网格 - 加上立体效果和可视范围区分
    for j in range(width):
        for i in range(height):
            x = i * grid_size_display
            y = j * grid_size_display
            
            # 判断是否在可视范围内
            in_vision = (abs(j - agent_grid_x) <= 1 and abs(i - agent_grid_y) <= 1)
            
            if grid[j,i] == 1:
                if in_vision:
                    # 可视范围内的自由空间 - 改为深色（与OpenCV一致）
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(40,50,40)', stroke='none'))
                    # 立体边框效果：深色调
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(70,90,70)', stroke_width=4))  # 上边
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(70,90,70)', stroke_width=4))  # 左边
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(20,30,20)', stroke_width=4))  # 下边
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(20,30,20)', stroke_width=4))  # 右边
                else:
                    # 不可视范围内的自由空间 - 深蓝灰色（与OpenCV一致）
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(30,40,50)', stroke='none'))
                    # 立体边框效果
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(50,70,90)', stroke_width=4))  # 上边
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(50,70,90)', stroke_width=4))  # 左边
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(15,25,35)', stroke_width=4))  # 下边
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(15,25,35)', stroke_width=4))  # 右边
            else:
                if in_vision:
                    # 可视范围内的障碍物 - 改为浅米色（与OpenCV一致）
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(220,220,200)', stroke='none'))
                    # 立体边框效果
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(250,250,230)', stroke_width=4))  # 上边
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(250,250,230)', stroke_width=4))  # 左边
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(190,190,170)', stroke_width=4))  # 下边
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(190,190,170)', stroke_width=4))  # 右边
                    # 添加浅色纹理点
                    dwg.add(dwg.circle(center=(x+12, y+12), r=2, fill='rgb(180,180,190)'))
                    dwg.add(dwg.circle(center=(x+28, y+28), r=2, fill='rgb(180,180,190)'))
                else:
                    # 不可视范围内的障碍物 - 改为浅灰色（与OpenCV一致）
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(200,200,210)', stroke='none'))
                    # 立体边框效果
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(230,230,240)', stroke_width=4))  # 上边
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(230,230,240)', stroke_width=4))  # 左边
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(170,170,180)', stroke_width=4))  # 下边
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(170,170,180)', stroke_width=4))  # 右边
                    # 添加浅色纹理点
                    dwg.add(dwg.circle(center=(x+12, y+12), r=2, fill='rgb(180,180,190)'))
                    dwg.add(dwg.circle(center=(x+28, y+28), r=2, fill='rgb(180,180,190)'))
    
    # 绘制轨迹路径 - 保持原有的渐变效果
    if len(trajectory) > 1:
        complete_trajectory = [init_pos] + trajectory
        
        # 绘制路径阴影
        for i in range(len(complete_trajectory)-1):
            dwg.add(dwg.line(start=(complete_trajectory[i][1]+2, complete_trajectory[i][0]+2), 
                           end=(complete_trajectory[i+1][1]+2, complete_trajectory[i+1][0]+2),
                           stroke='rgb(40,40,40)', stroke_width=8))
        
        # 绘制路径主线条
        for i in range(len(complete_trajectory)-1):
            progress = i / max(1, len(complete_trajectory) - 2)
            green_r = int(0 + progress * 100)
            green_g = int(100 + progress * 155)
            green_b = int(0 + progress * 100)
            
            dwg.add(dwg.line(start=(complete_trajectory[i][1], complete_trajectory[i][0]), 
                           end=(complete_trajectory[i+1][1], complete_trajectory[i+1][0]),
                           stroke=f'rgb({green_r},{green_g},{green_b})', stroke_width=6))
        
        # 移除路径中心亮线（根据用户要求）
        # for i in range(len(complete_trajectory)-1):
        #     dwg.add(dwg.line(start=(complete_trajectory[i][1], complete_trajectory[i][0]), 
        #                    end=(complete_trajectory[i+1][1], complete_trajectory[i+1][0]),
        #                    stroke='rgb(150,255,150)', stroke_width=2))
        
        # 绘制路径点
        for i, point in enumerate(complete_trajectory):
            progress = i / max(1, len(complete_trajectory) - 1)
            dot_r = int(50 + progress * 100)
            dot_g = int(150 + progress * 105)
            dot_b = int(50 + progress * 100)
            
            dwg.add(dwg.circle(center=(point[1], point[0]), r=4, 
                             fill=f'rgb({dot_r},{dot_g},{dot_b})'))
            dwg.add(dwg.circle(center=(point[1], point[0]), r=6, 
                             fill='none', stroke='rgb(100,200,100)', stroke_width=1))
    
    # 绘制起始位置标记 - 蓝色方框 (修正颜色BGR->RGB)
    init_x, init_y = init_pos[1], init_pos[0]
    dwg.add(dwg.rect(insert=(init_x-18, init_y-18), size=(36, 36), 
                   fill='none', stroke='rgb(0,120,180)', stroke_width=6))
    
    # 绘制目标位置 - 多层圆圈效果
    goal_x, goal_y = goal_pos[1], goal_pos[0]
    dwg.add(dwg.circle(center=(goal_x, goal_y), r=24, 
                     fill='none', stroke='rgb(50,200,50)', stroke_width=4))
    dwg.add(dwg.circle(center=(goal_x, goal_y), r=16, fill='rgb(0,180,0)'))
    dwg.add(dwg.circle(center=(goal_x-4, goal_y-4), r=4, fill='rgb(100,255,100)'))
    
    # 绘制可视范围边界指示器
    vision_x1 = (agent_grid_y - 1) * grid_size_display
    vision_y1 = (agent_grid_x - 1) * grid_size_display
    vision_x2 = (agent_grid_y + 2) * grid_size_display
    vision_y2 = (agent_grid_x + 2) * grid_size_display
    
    # 可视范围边框
    dwg.add(dwg.rect(insert=(vision_x1, vision_y1), 
                   size=(vision_x2-vision_x1, vision_y2-vision_y1),
                   fill='none', stroke='rgb(150,255,180)', stroke_width=3))
    
    # 绘制agent三角形 - 完整的视觉效果
    agent_x, agent_y = agent_pos[1], agent_pos[0]
    triangle_size = 28
    
    # 根据动作确定三角形朝向
    if agent_action == 0:  # 向右
        points = [(agent_x - triangle_size//2, agent_y - triangle_size//2),
                 (agent_x - triangle_size//2, agent_y + triangle_size//2),
                 (agent_x + triangle_size//2, agent_y)]
        inner_points = [(agent_x - (triangle_size-8)//2, agent_y - (triangle_size-8)//2),
                       (agent_x - (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x + (triangle_size-8)//2, agent_y)]
    elif agent_action == 1:  # 向左
        points = [(agent_x + triangle_size//2, agent_y - triangle_size//2),
                 (agent_x + triangle_size//2, agent_y + triangle_size//2),
                 (agent_x - triangle_size//2, agent_y)]
        inner_points = [(agent_x + (triangle_size-8)//2, agent_y - (triangle_size-8)//2),
                       (agent_x + (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x - (triangle_size-8)//2, agent_y)]
    elif agent_action == 2:  # 向下
        points = [(agent_x - triangle_size//2, agent_y - triangle_size//2),
                 (agent_x + triangle_size//2, agent_y - triangle_size//2),
                 (agent_x, agent_y + triangle_size//2)]
        inner_points = [(agent_x - (triangle_size-8)//2, agent_y - (triangle_size-8)//2),
                       (agent_x + (triangle_size-8)//2, agent_y - (triangle_size-8)//2),
                       (agent_x, agent_y + (triangle_size-8)//2)]
    elif agent_action == 3:  # 向上
        points = [(agent_x - triangle_size//2, agent_y + triangle_size//2),
                 (agent_x + triangle_size//2, agent_y + triangle_size//2),
                 (agent_x, agent_y - triangle_size//2)]
        inner_points = [(agent_x - (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x + (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x, agent_y - (triangle_size-8)//2)]
    else:  # 默认向上
        points = [(agent_x - triangle_size//2, agent_y + triangle_size//2),
                 (agent_x + triangle_size//2, agent_y + triangle_size//2),
                 (agent_x, agent_y - triangle_size//2)]
        inner_points = [(agent_x - (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x + (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x, agent_y - (triangle_size-8)//2)]
    
    # 绘制三角形阴影 - 让阴影比三角形稍大一些
    shadow_size = triangle_size + 4  # 阴影比主体大4px
    if agent_action == 0:  # 向右
        shadow_points = [(agent_x - shadow_size//2 + 4, agent_y - shadow_size//2 + 4),
                        (agent_x - shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x + shadow_size//2 + 4, agent_y + 4)]
    elif agent_action == 1:  # 向左
        shadow_points = [(agent_x + shadow_size//2 + 4, agent_y - shadow_size//2 + 4),
                        (agent_x + shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x - shadow_size//2 + 4, agent_y + 4)]
    elif agent_action == 2:  # 向下
        shadow_points = [(agent_x - shadow_size//2 + 4, agent_y - shadow_size//2 + 4),
                        (agent_x + shadow_size//2 + 4, agent_y - shadow_size//2 + 4),
                        (agent_x + 4, agent_y + shadow_size//2 + 4)]
    elif agent_action == 3:  # 向上
        shadow_points = [(agent_x - shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x + shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x + 4, agent_y - shadow_size//2 + 4)]
    else:  # 默认向上
        shadow_points = [(agent_x - shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x + shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x + 4, agent_y - shadow_size//2 + 4)]
    
    dwg.add(dwg.polygon(points=shadow_points, fill='rgb(30,30,30)'))
    
    # 绘制三角形主体 (修正颜色：OpenCV是BGR，SVG是RGB)
    dwg.add(dwg.polygon(points=points, fill='rgb(255,0,0)', stroke='rgb(150,0,0)', stroke_width=6))
    
    # 绘制内部高光三角形
    dwg.add(dwg.polygon(points=inner_points, fill='rgb(255,100,100)', stroke='rgb(255,150,150)', stroke_width=2))
    
    # 绘制中心亮点
    dwg.add(dwg.circle(center=(agent_x, agent_y), r=4, fill='rgb(255,200,200)'))
    dwg.add(dwg.circle(center=(agent_x, agent_y), r=2, fill='rgb(255,255,255)'))
    
    # ============ 扰动箭头已移除（SVG版本）============
    # 移除扰动箭头功能
    
    # 保存SVG文件
    dwg.save()
    print(f"SVG snapshot saved to {pth}")

def render(grid, state, goal, valid = True, show_vision_mask = False):
        
        state_x = int(state[0])
        state_y = int(state[1])

        food_x = int(goal[0])
        food_y = int(goal[1])

        grid_size_display = 40  # 提升到4倍分辨率 (2x2)
        width, height = grid.shape[0], grid.shape[1]
        img = np.zeros((width * grid_size_display, height * grid_size_display, 3), np.uint8)
        
        for j in range(width):
            for i in range(height):
                if grid[j,i] == 1:
                    # 自由空间 - 改为深色背景
                    cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), 
                                (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), 
                                (30, 40, 50), -1)  # 深蓝灰色
                    
                    # 立体边框效果：深色调
                    # 上边框 - 深色
                    cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                           ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                           (50, 70, 90), 4)
                    # 左边框 - 深色
                    cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                           (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                           (50, 70, 90), 4)
                    # 下边框 - 更深色
                    cv2.line(img, (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                           ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                           (15, 25, 35), 4)
                    # 右边框 - 更深色
                    cv2.line(img, ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                           ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                           (15, 25, 35), 4)
                else:
                    # 障碍物 - 改为浅色效果
                    cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), 
                                (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), 
                                (200, 200, 210), -1)  # 浅灰色基底
                    
                    # 石墙立体边框效果 (浅色调)
                    cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                           ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                           (230, 230, 240), 4)
                    cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                           (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                           (230, 230, 240), 4)
                    cv2.line(img, (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                           ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                           (170, 170, 180), 4)
                    cv2.line(img, ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                           ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                           (170, 170, 180), 4)
                    
                    # 添加浅色纹理效果
                    cv2.circle(img, (i * grid_size_display + 12, j * grid_size_display + 12), 2, (180, 180, 190), -1)
                    cv2.circle(img, (i * grid_size_display + 28, j * grid_size_display + 28), 2, (180, 180, 190), -1)
                
                # 智能体位置不在这里绘制，稍后在最上层绘制三角形
                pass
        
        # 目标位置 - 改进的绿色圆形，添加脉冲效果 (放大)
        # 外圈光晕
        cv2.circle(img, (food_y * grid_size_display + grid_size_display//2, food_x * grid_size_display + grid_size_display//2), 
                 24, (50, 200, 50), 4, cv2.LINE_AA)
        # 主体
        cv2.circle(img, (food_y * grid_size_display + grid_size_display//2, food_x * grid_size_display + grid_size_display//2), 
                 16, (0, 180, 0), -1, cv2.LINE_AA)
        # 内部高光
        cv2.circle(img, (food_y * grid_size_display + grid_size_display//2 - 4, food_x * grid_size_display + grid_size_display//2 - 4), 
                 4, (100, 255, 100), -1, cv2.LINE_AA)

        if not valid:
            # draw a big red cross (放大)
            cv2.line(img, (0, 0), (img.shape[1], img.shape[0]), (0, 0, 255), 10, cv2.LINE_AA)
            cv2.line(img, (0, img.shape[0]), (img.shape[1], 0), (0, 0, 255), 10, cv2.LINE_AA)
            cv2.putText(img, "invalid map", (int(img.shape[1]/2) - 200, int(img.shape[0]/2)), cv2.FONT_HERSHEY_SIMPLEX, 2.0, (200, 200, 200), 6, cv2.LINE_AA)
            cv2.putText(img, "invalid map", (int(img.shape[1]/2) - 200, int(img.shape[0]/2)), cv2.FONT_HERSHEY_SIMPLEX, 2.0, (200, 0, 0), 4, cv2.LINE_AA)
        return img

event_type = ""
event_x = 0
event_y = 0

def run_editor(landscape, state, goal, map_size=12):

    global event_type, event_x, event_y

    # define mouse callback function
    def input_cb(event,x,y,flags,param):
        global event_type, event_x, event_y
        if event == cv2.EVENT_MOUSEWHEEL:
            event_x = x
            event_y = y
            event_type = "flip_space"
        elif event == cv2.EVENT_LBUTTONUP:
            event_x = x
            event_y = y
            event_type = "set_start"
        elif event == cv2.EVENT_RBUTTONUP:
            event_x = x
            event_y = y
            event_type = "set_goal"

    cv2.namedWindow("img", cv2.WINDOW_GUI_NORMAL)
    cv2.setMouseCallback("img", input_cb)

    grid = np.array(landscape).reshape(map_size, map_size).transpose()
    valid = True
    grid_size_display = 20

    while True:

        if event_type == "flip_space":
            grid[event_y//grid_size_display, event_x//grid_size_display] = 1 - grid[event_y//grid_size_display, event_x//grid_size_display]
            event_type = "flip_space_done"
        elif event_type == "set_start":
            state = [event_y//grid_size_display, event_x//grid_size_display]
            event_type = "set_start_done"
        elif event_type == "set_goal":
            goal = [event_y//grid_size_display, event_x//grid_size_display]
            event_type = "set_goal_done"

        num_labels, labels, stats, centroids, num_freespace, landscape_img = check_num_labels(grid, map_size, map_size)
        non_zeros = np.count_nonzero(landscape)
        if not (num_labels == 2 and non_zeros >= 5):
            valid = False
        else:
            valid = True
        
        img = render(grid, state, goal, valid)
        cv2.imshow("img", img)
        k = cv2.waitKey(1)
        if k == ord('q'):
            break
        elif k == ord('s'):
            pth = "./logs/test.json"
            grid0 = grid.transpose()
            grid1 = grid0.reshape(map_size*map_size)
            save_current_task([grid1], state[0], state[1], goal[0], goal[1], pth)
            print("task saved to {}".format(pth))
        elif k == ord('r'):
            grid = np.array(landscape).reshape(map_size, map_size).transpose()
    
    grid_ = grid.transpose()
    grid_ = grid_.reshape(map_size*map_size).tolist()
    return grid_


def main():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration

    k1 = jax.random.PRNGKey(npr.randint(0, 1000000))

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    """ create landscape
    """
    random_task = True
    # check if file on rpl_config.task_pth exists
    if os.path.isfile(rpl_config.task_pth):
        random_task = False

    if random_task:
        landscape = generate_maze_pool(num_mazes=1, width=10, height=10)
        landscape = padding_landscapes(landscape, width=12, height=12)
    else:
        landscape, state, goal = load_task(pth = rpl_config.task_pth)
        landscape = [landscape]

    print("landscape :")
    print(landscape)

    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    """ create grid env
    """
    start_time = time.time()
    GE = GridEnv(landscapes = landscape, width = 12, height = 12, num_envs_per_landscape = 1)
    GE.reset()
    print("time taken to create envs: ", time.time() - start_time)

    if not random_task:
        # set states of GE
        GE.batched_states = GE.batched_states.at[0, 0].set(state[0])
        GE.batched_states = GE.batched_states.at[0, 1].set(state[1])
        # set goals of GE
        GE.batched_goals = GE.batched_goals.at[0, 0].set(goal[0])
        GE.batched_goals = GE.batched_goals.at[0, 1].set(goal[1])
        GE.init_batched_states, GE.init_batched_goals = jnp.copy(GE.batched_states), jnp.copy(GE.batched_goals)
        GE.batched_goal_reached = batch_compute_goal_reached(GE.batched_states, GE.batched_goals)
        GE.last_batched_goal_reached = jnp.copy(GE.batched_goal_reached)
        GE.concat_obs = get_ideal_obs_vmap(GE.batched_envs, GE.batched_states, GE.batched_goals, GE.last_batched_goal_reached)

    concat_obs = GE.concat_obs

    rnn_state = model.initial_state(GE.num_envs)

    step_count = 0
    render_id = 0

    trajectory = []
    rnn_state_waterfall = []

    reset_ = True

    step_by_step = False
    manual_action = 0
    
    # 扰动可视化相关变量
    perturbation_arrows = []  # 存储扰动箭头信息 [(x, y, action, timestamp), ...]
    perturbation_display_duration = 3000  # 箭头显示时长（毫秒）
    
    # 神经状态序列记录变量
    current_cycle_neural_states = []  # 记录当前周期的神经状态序列
    
    while True:

        step_count += 1

        """ model forward and step the env
        """
        rnn_state, y1 = model_forward(params, rnn_state, concat_obs, model)
        
        if not step_by_step:
            batched_actions = get_action_vmap(y1)
        else:
            batched_actions = jnp.array([manual_action])

        batched_goal_reached, concat_obs = GE.step(batched_actions, reset = reset_)

        rnn_state_waterfall.append(rnn_state[0].tolist())
        
        # 记录当前周期的神经状态序列
        current_cycle_neural_states.append(rnn_state[render_id].tolist())

        # 在执行动作前记录当前位置用于轨迹绘制 (适应4倍分辨率)
        trajectory.append([40 * GE.batched_states[render_id][0]+20, 40 * GE.batched_states[render_id][1]+20])

        if batched_goal_reached[render_id]:
            # Agent到达终点时保存SVG快照
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            svg_filename = f"./logs/navigation_snapshot_{timestamp}.svg"
            
            # 准备SVG保存所需的数据 - 确保转换为Python原生类型
            current_grid = np.array(GE.batched_envs[render_id])
            init_pos = [int(40 * int(GE.init_batched_states[render_id][0]) + 20), 
                       int(40 * int(GE.init_batched_states[render_id][1]) + 20)]
            goal_pos = [int(40 * int(GE.batched_goals[render_id][0]) + 20), 
                       int(40 * int(GE.batched_goals[render_id][1]) + 20)]
            agent_pos = [int(40 * int(GE.batched_states[render_id][0]) + 20), 
                        int(40 * int(GE.batched_states[render_id][1]) + 20)]
            current_action = int(batched_actions[render_id]) if len(batched_actions) > render_id else 0
            
            # 转换轨迹数据为Python原生类型
            trajectory_native = [[int(point[0]), int(point[1])] for point in trajectory]
            
            # 保存SVG快照（传递当前的扰动箭头信息）
            save_svg_snapshot(current_grid, trajectory_native, init_pos, goal_pos, agent_pos, current_action, perturbation_arrows, svg_filename)
            
            # ============ 保存当前周期的神经状态序列SVG ============
            if len(current_cycle_neural_states) > 0:
                neural_states_svg_path = "./logs/neural_states_sequence.svg"
                save_neural_states_sequence_svg(current_cycle_neural_states, neural_states_svg_path)
                print(f"Saved neural states sequence with {len(current_cycle_neural_states)} states")
                
                # ============ 生成并显示OpenCV预览 ============
                sequence_img = create_neural_states_sequence_opencv(current_cycle_neural_states)
                if sequence_img is not None:
                    cv2.imshow("Neural States Sequence", sequence_img)
                
                # 清空当前周期的神经状态序列，准备下一个周期
                current_cycle_neural_states.clear()
            
            # 清除扰动箭头（每轮保存后立即清除）
            perturbation_arrows.clear()
            
            trajectory.clear()

        """ render the env
        """
        if rpl_config.visualization == "True" or rpl_config.video_output == "True":
            # 传递当前动作信息给render函数，用于绘制朝向正确的agent三角形
            current_action = batched_actions[render_id] if len(batched_actions) > render_id else 0
            img = GE.render(env_id = render_id, current_action = current_action)
            if len(trajectory) > 1:
                # 创建完整的路径序列，确保包含起始位置
                complete_trajectory = []
                
                # 强制插入起始位置（初始位置）
                init_x = int(GE.init_batched_states[render_id][0])
                init_y = int(GE.init_batched_states[render_id][1])
                start_pos = [40 * init_x + 20, 40 * init_y + 20]
                complete_trajectory.append(start_pos)
                
                # 添加所有已记录的轨迹点
                complete_trajectory.extend(trajectory)
                
                # 改进的路径绘制 - 更有质感的效果
                for i in range(len(complete_trajectory)-1):
                    # 计算渐变颜色 - 从深绿到亮绿
                    progress = i / max(1, len(trajectory) - 2) if len(trajectory) > 2 else 0
                    
                    # 渐变绿色：从深绿(0,100,0)到亮绿(100,255,100)
                    green_r = int(0 + progress * 100)
                    green_g = int(100 + progress * 155)
                    green_b = int(0 + progress * 100)
                    
                    # 1. 先绘制阴影效果 (偏移的深色线条)
                    cv2.line(img, 
                            (int(complete_trajectory[i][1]) + 2, int(complete_trajectory[i][0]) + 2), 
                            (int(complete_trajectory[i+1][1]) + 2, int(complete_trajectory[i+1][0]) + 2), 
                            (40, 40, 40), 8, cv2.LINE_AA)
                    
                    # 2. 绘制主路径线条 (较粗的渐变色)
                    cv2.line(img, 
                            (int(complete_trajectory[i][1]), int(complete_trajectory[i][0])), 
                            (int(complete_trajectory[i+1][1]), int(complete_trajectory[i+1][0])), 
                            (green_r, green_g, green_b), 6, cv2.LINE_AA)
                    
                    # 3. 绘制路径中心的亮线 (更细的亮色线条)
                    cv2.line(img, 
                            (int(complete_trajectory[i][1]), int(complete_trajectory[i][0])), 
                            (int(complete_trajectory[i+1][1]), int(complete_trajectory[i+1][0])), 
                            (150, 255, 150), 2, cv2.LINE_AA)
                
                # 4. 在路径点上绘制小圆点标记
                for i in range(len(complete_trajectory)):
                    progress = i / max(1, len(complete_trajectory) - 1)
                    dot_r = int(50 + progress * 100)
                    dot_g = int(150 + progress * 105)
                    dot_b = int(50 + progress * 100)
                    
                    # 绘制路径点的小圆圈
                    cv2.circle(img, (int(complete_trajectory[i][1]), int(complete_trajectory[i][0])), 
                             4, (dot_r, dot_g, dot_b), -1, cv2.LINE_AA)
                    cv2.circle(img, (int(complete_trajectory[i][1]), int(complete_trajectory[i][0])), 
                             6, (100, 200, 100), 1, cv2.LINE_AA)
            
            # ============ 在路径绘制完成后再次绘制Agent三角形，确保不被遮挡 ============
            state_x = int(GE.batched_states[render_id][0])
            state_y = int(GE.batched_states[render_id][1])
            center_x = state_y * 40 + 20  # 40是grid_size_display
            center_y = state_x * 40 + 20
            triangle_size = 28
            
            # 根据动作确定三角形朝向
            if current_action == 0:  # 向右
                triangle_pts = np.array([
                    [center_x - triangle_size//2, center_y - triangle_size//2],
                    [center_x - triangle_size//2, center_y + triangle_size//2],
                    [center_x + triangle_size//2, center_y]
                ], np.int32)
            elif current_action == 1:  # 向左
                triangle_pts = np.array([
                    [center_x + triangle_size//2, center_y - triangle_size//2],
                    [center_x + triangle_size//2, center_y + triangle_size//2],
                    [center_x - triangle_size//2, center_y]
                ], np.int32)
            elif current_action == 2:  # 向下
                triangle_pts = np.array([
                    [center_x - triangle_size//2, center_y - triangle_size//2],
                    [center_x + triangle_size//2, center_y - triangle_size//2],
                    [center_x, center_y + triangle_size//2]
                ], np.int32)
            elif current_action == 3:  # 向上
                triangle_pts = np.array([
                    [center_x - triangle_size//2, center_y + triangle_size//2],
                    [center_x + triangle_size//2, center_y + triangle_size//2],
                    [center_x, center_y - triangle_size//2]
                ], np.int32)
            else:  # 默认向上
                triangle_pts = np.array([
                    [center_x - triangle_size//2, center_y + triangle_size//2],
                    [center_x + triangle_size//2, center_y + triangle_size//2],
                    [center_x, center_y - triangle_size//2]
                ], np.int32)
            
            # 绘制三角形阴影
            shadow_pts = triangle_pts + 4
            cv2.fillPoly(img, [shadow_pts], (30, 30, 30))
            
            # 绘制三角形外轮廓
            cv2.polylines(img, [triangle_pts], True, (0, 0, 150), 6, cv2.LINE_AA)
            
            # 绘制三角形主体
            cv2.fillPoly(img, [triangle_pts], (0, 0, 255))
            
            # 绘制内部高光
            inner_size = triangle_size - 8
            if current_action == 0:  # 向右
                inner_triangle_pts = np.array([
                    [center_x - inner_size//2, center_y - inner_size//2],
                    [center_x - inner_size//2, center_y + inner_size//2],
                    [center_x + inner_size//2, center_y]
                ], np.int32)
            elif current_action == 1:  # 向左
                inner_triangle_pts = np.array([
                    [center_x + inner_size//2, center_y - inner_size//2],
                    [center_x + inner_size//2, center_y + inner_size//2],
                    [center_x - inner_size//2, center_y]
                ], np.int32)
            elif current_action == 2:  # 向下
                inner_triangle_pts = np.array([
                    [center_x - inner_size//2, center_y - inner_size//2],
                    [center_x + inner_size//2, center_y - inner_size//2],
                    [center_x, center_y + inner_size//2]
                ], np.int32)
            elif current_action == 3:  # 向上
                inner_triangle_pts = np.array([
                    [center_x - inner_size//2, center_y + inner_size//2],
                    [center_x + inner_size//2, center_y + inner_size//2],
                    [center_x, center_y - inner_size//2]
                ], np.int32)
            else:  # 默认向上
                inner_triangle_pts = np.array([
                    [center_x - inner_size//2, center_y + inner_size//2],
                    [center_x + inner_size//2, center_y + inner_size//2],
                    [center_x, center_y - inner_size//2]
                ], np.int32)
            
            cv2.fillPoly(img, [inner_triangle_pts], (100, 100, 255))
            cv2.polylines(img, [inner_triangle_pts], True, (150, 150, 255), 2, cv2.LINE_AA)
            cv2.circle(img, (center_x, center_y), 4, (200, 200, 255), -1, cv2.LINE_AA)
            cv2.circle(img, (center_x, center_y), 2, (255, 255, 255), -1, cv2.LINE_AA)
            
            # ============ 扰动箭头已移除 ============
            # 移除扰动箭头显示功能

        """ scene display
        """
        if rpl_config.visualization == "True":
            
            # ============ 实时可视化神经网络隐状态 ============
            # 获取当前隐状态并可视化
            current_hidden_state = rnn_state[render_id]  # 获取当前环境的隐状态
            visualize_hidden_state(current_hidden_state, "Neural Hidden State Visualization")

            cv2.imshow("img", img)

            if step_by_step:
                k = cv2.waitKey(0)
            else:
                k = cv2.waitKey(1)
            if k == ord('r'): 
                rnn_state_waterfall.clear()
                current_cycle_neural_states.clear()  # 清空神经状态序列
                rnn_state = model.initial_state(GE.num_envs)
                GE.rnd_goal_collection = get_rnd_goal_collection_vmap(GE.env_keys, GE.batched_envs, GE.width, GE.height, GE.num_free_spaces)
                GE.reset()
                if not random_task:
                    # set states of GE
                    GE.batched_states = GE.batched_states.at[0, 0].set(state[0])
                    GE.batched_states = GE.batched_states.at[0, 1].set(state[1])
                    # set goals of GE
                    GE.batched_goals = GE.batched_goals.at[0, 0].set(goal[0])
                    GE.batched_goals = GE.batched_goals.at[0, 1].set(goal[1])
                    GE.init_batched_states, GE.init_batched_goals = jnp.copy(GE.batched_states), jnp.copy(GE.batched_goals)
                    GE.batched_goal_reached = batch_compute_goal_reached(GE.batched_states, GE.batched_goals)
                    GE.last_batched_goal_reached = jnp.copy(GE.batched_goal_reached)
                    GE.concat_obs = get_ideal_obs_vmap(GE.batched_envs, GE.batched_states, GE.batched_goals, GE.last_batched_goal_reached)
                    concat_obs = GE.concat_obs
                    
                trajectory.clear()
            elif k == ord('p'):
                # 1. 记录扰动发生前的位置
                perturbation_x = int(GE.batched_states[render_id][0])
                perturbation_y = int(GE.batched_states[render_id][1])
                
                # 2. take a random action
                k2, _ = jax.random.split(k1)
                k1 = k2
                random_action = jax.random.randint(k1, shape=(GE.num_envs, ), minval=0, maxval=4)
                action_value = int(random_action[render_id])
                
                # 3. 记录扰动箭头信息 (位置, 动作, 时间戳)
                current_time = time.time() * 1000  # 转换为毫秒
                perturbation_arrows.append((perturbation_x, perturbation_y, action_value, current_time))
                
                batched_goal_reached, concat_obs = GE.step(random_action)
                print("random action: ", random_action, f"at position ({perturbation_x}, {perturbation_y})")
                
                # 4. 更新轨迹记录，添加随机动作后的新位置 (适应4倍分辨率)
                trajectory.append([40 * GE.batched_states[render_id][0]+20, 40 * GE.batched_states[render_id][1]+20])
                
                # 5. 检查是否到达目标，如果是则清空轨迹
                if batched_goal_reached[render_id]:
                    trajectory.clear()

            elif k == ord('n'):
                if random_task:
                    rnn_state = model.initial_state(GE.num_envs)
                    landscape = generate_maze_pool(num_mazes=1, width=10, height=10)
                    landscape = padding_landscapes(landscape, width=12, height=12)
                    GE.set_landscapes(landscape)
                    GE.reset()
                    trajectory.clear()
                    rnn_state_waterfall.clear()
                    current_cycle_neural_states.clear()  # 清空神经状态序列
            elif k == ord('q'):
                break
            elif k == ord('t'):
                step_by_step = not step_by_step
            elif k == ord('w'):
                manual_action = 3
            elif k == ord('s'):
                manual_action = 2
            elif k == ord('a'):
                manual_action = 1
            elif k == ord('d'):
                manual_action = 0
            elif k == ord('e'):
                new_landscape = run_editor(landscape[0], GE.batched_states[0], GE.batched_goals[0])
                GE.set_landscapes_only([new_landscape])
                landscape[0] = new_landscape

if __name__ == "__main__":
    main()
